#ifndef TrtEngine_H
#define TrtEngine_H
#include <NvInfer.h>
#include <NvOnnxParser.h>
#include <cuda_runtime_api.h>
#include <iostream>
#include <fstream>
#include <sstream>
#include <memory>
#include <opencv2/opencv.hpp>
#include "yolov5.h"
using namespace std;
/**
  定义TensorRT的工具类
**/
static class Logger : public nvinfer1::ILogger
{
public:
    void log(Severity severity, const char *msg) noexcept
    {
        if ((severity == Severity::kERROR) || (severity == Severity::kINTERNAL_ERROR))
        {
            cout << msg << endl;
        }
    }
} gLogger;
// static inline int read_files_in_dir(const char *p_dir_name, std::vector<std::string> &file_names) {
//     DIR *p_dir = opendir(p_dir_name);
//     if (p_dir == nullptr) {
//         return -1;
//     }

//    struct dirent* p_file = nullptr;
//    while ((p_file = readdir(p_dir)) != nullptr) {
//        if (strcmp(p_file->d_name, ".") != 0 &&
//            strcmp(p_file->d_name, "..") != 0) {
//            //std::string cur_file_name(p_dir_name);
//            //cur_file_name += "/";
//            //cur_file_name += p_file->d_name;
//            std::string cur_file_name(p_file->d_name);
//            file_names.push_back(cur_file_name);
//        }
//    }

//    closedir(p_dir);
//    return 0;
//}

#define CHECK(status)                              \
    do                                             \
    {                                              \
        cudaError_t err = (status);                \
        if (err != cudaSuccess)                    \
        {                                          \
            fprintf(stderr, "API error"            \
                            "%s:%d Returned:%d\n", \
                    __FILE__, __LINE__, err);      \
            exit(1);                               \
        }                                          \
    } while (0)
class TrtEngine
{
private:
    Yolov5 yolov5;
    int netHeight = 1280; //网络的输入高度
    int netWidth = 1280;  //网络的输出高度

    int batch_size, channel, inputHeight, intputWidth;
    int output_size;
    void *buffers[2] = {nullptr};

public:
    int output_width = 7; //网络的classes总数
    string path;
    float boxThreshold, nmsThreshold;
    TrtEngine();
    void init();
    uint8_t *d_img = {nullptr};
    ~TrtEngine(){};
    void loadOnnxEngine(const string onnx_filename);         //加载onnx模型并转换为trt模型
    nvinfer1::ICudaEngine *loadTRTEngine(string enginePath); //加载trt模型
    cv::Mat preprocess_img(cv::Mat &img);
    cv::Mat letterbox(cv::Mat &img);   //实现yolov5的letterbox图片预处理
    float *ProcessImage(cv::Mat &img); //获得trt执行结果
    void mat2float(float *data, cv::Mat img);
    void mat2float2(float *data, cv::Mat img);
    cv::Rect get_rect(cv::Mat &img, float bbox[4]);
    nvinfer1::ICudaEngine *trtengine;
    nvinfer1::IExecutionContext *mp_context;
    cudaStream_t stream;
};

#endif // TrtEngine_H
